import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import DBSCAN
from lib.calcPoints import calcPoints
import json
import cv2
from matplotlib.backends.backend_agg import FigureCanvasAgg
import matplotlib
import math

# matplotlib.use('agg')


def matplotlib_to_opencv(matplotlib_img):
    """
    ��matplotlibͼ��ת��ΪOpenCVͼ��
    """
    # ��matplotlibͼ��ת��ΪRGB��ʽ
    rgb_image = matplotlib_img[:, :, :3].copy()
    # ��RGBת��ΪBGR
    bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
    return bgr_image


def getCenters(
    points, validDistance, eps=0.25, min_samples=5, debug=False, size=(100, 100)
):
    points = np.array(points)

    dbscan = DBSCAN(eps=eps, min_samples=min_samples)
    # ���ģ�Ͳ�Ԥ������ǩ
    labels = dbscan.fit_predict(points)

    dic = {}
    # ����ͬ��ɫ�ĵ������
    centers = []
    for i in range(len(labels)):
        if labels[i] not in dic:
            dic[labels[i]] = []
        dic[labels[i]].append(points[i])
        if labels[i] == -1:
            continue

    for i in dic:
        c = np.mean(dic[i], axis=0)
        if math.sqrt(c[0] ** 2 + c[1] ** 2) > validDistance:
            continue
        centers.append(c)

    centers = np.array(centers)
    opencv_img = None

    if debug:
        fig, ax = plt.subplots()
        # 画幅大小
        ax.set_xlim(-size[0], size[1])
        ax.set_ylim(0, size[1])

        ax.plot(0, 0, "ko", markersize=2)

        ax.scatter(points[:, 0], points[:, 1], c=labels, cmap="tab10")
        for i in range(len(centers)):
            if centers[i][1] < 30:
                continue
            ax.scatter(centers[i][0], centers[i][1], c="red", marker="x", s=100)

        canvas = FigureCanvasAgg(fig)
        canvas.draw()

        buf = canvas.buffer_rgba()
        height, width, channels = buf.shape

        opencv_img = np.frombuffer(buf, dtype=np.uint8).reshape(height, width, channels)

    return centers, opencv_img


if __name__ == "__main__":
    import rospy, time
    from sensor_msgs.msg import LaserScan

    rospy.init_node("lsidar_test_scan")

    while True:
        t = time.time()
        r = rospy.wait_for_message("/scan", LaserScan)
        X = np.array(calcPoints(r.ranges[:1000]))
        cs, img = getCenters(X, debug=0)
        # cv2.imshow("DEBUG_IMG", img)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
        # fps
        print(1 / (time.time() - t))

    cv2.destroyAllWindows()
